import cPickle as pickle
from numpy import linspace, complex128, append, zeros, sin, pi
from utils import stoc_eval, traj_fid

def fid_append(var_tpl, var_names, gammas, gamma_names, pulse_f, input_dict):
    """
    returns a dictionary consisting of the input variables in var_tpl,
    with the names given in var_names, as well as a measurement
    fidelity obtained by evaluation until 
    """
    
    #Eliminating nasty hack from fidelity_grid.py
    #dep_dict = {name: var for name, var in zip(var_names, var_tpl)}
    
    #Getting rid of dict comprehension to appease old versions of 
    #Python on the Juelich cluster
    dep_dict = {}
    for name, var in zip(var_names, var_tpl):
        dep_dict[name] = var
    
    #new nasty hack
    meas_t = 10.
    
    dt = input_dict['dt']
    T = meas_t
    N = int(round(T / dt))
    tvec = linspace(0., T, N)
    
    epsilon = pulse_f(tvec, *var_tpl)
    
    #Extreme Ugliness: add 2/chi to the time vector here, 
    #along with the requisite number of zeros to epsilon.

    T += 2 #Simulation time gets an extra 2/chi
    N = int(round(T / dt))
    tvec = linspace(0., T, N)

    zeros_len = len(tvec) - len(epsilon)
    epsilon = append(epsilon, zeros(zeros_len))

    input_dict['T'] = T
    input_dict['tvec'] = tvec
    input_dict['epsilon'] = epsilon
    input_dict['N'] = N
  
    gamma_d = gammas[0]
    gamma_r = gammas[1]
    Ld1 = input_dict['Ld1']; Ld2 = input_dict['Ld2']
    Ld3 = input_dict['Ld3']
    Lr1 = input_dict['Lr1']; Lr2 = input_dict['Lr2']
    Lr3 = input_dict['Lr3']
    kappa1 = input_dict['kappa1']; kappa2 = input_dict['kappa2']
    #Hard-coded Purcell decay factor 1/1600 
    D1 = (0.5 * gamma_d * Ld1 + 0.5 * gamma_d * Ld2 + 0.5 * gamma_d * Ld3 + 
            0 * (1./1600. * (kappa1 + kappa2) + gamma_r ) * (Lr1 + Lr2 + Lr3)).astype(complex128)

    input_dict.update({'D1': D1, 'eta': 0.5, 'tmeas_on': 0.,
                'tmeas_off': meas_t, 'epsilon': epsilon, 'N': N, 'T': meas_t, 'tvec':tvec})
    
    iter_check = True
    tol = 5*10**-4
    maxiter = 10**5
    if iter_check:
        fid, ci, k = stoc_eval(traj_fid, (input_dict,'ideal'), tol, maxiter=maxiter, iter_check=True)
        dep_dict.update({'fid':fid, 'confidence':ci, 'iters':k})    
    else:
        fid, ci = stoc_eval(traj_fid, (input_dict,'ideal'), tol, maxiter=maxiter, iter_check=False)
        dep_dict.update({'fid':fid, 'confidence':ci})    
    
    return dep_dict

def pulse_from_file(total_height, filename):
    #that file better just be a pickled array    
    phil = open(filename, 'r')
    try:
        pulse = pickle.load(phil)
    finally:
        phil.close()
    pulse = pulse/max(pulse) * total_height
    return pulse

def pulse_from_file_jr(tvec, total_height, filename):
    #that file better just be a pickled array    
    phil = open(filename, 'r')
    try:
        pulse = pickle.load(phil)
    finally:
        phil.close()

    pulse_time = np.linspace(tvec[0], tvec[-1], len(pulse))
    interp_pulse = np.interp(tvec, pulse_time, pulse)

    pulse = interp_pulse/max(interp_pulse) * total_height
    return pulse

def sine_pulse(tvec, amplitude, meas_t):
    return amplitude * sin(tvec * pi / meas_t)

def raised_sine_pulse(tvec, amplitude, y_offset, meas_t):
    if (y_offset < 0.) or (y_offset > 1.):
        raise ValueError("y_offset is supposed to be between 0 and 1.")
    return amplitude * ( y_offset + (1. - y_offset) * sin(tvec * pi / meas_t))

if __name__ == "__main__":
    from utils import prep_dict_indep, current_timestamp
    import numpy as np
    from sys import argv
    master_dir = '/usr/users/iff_th3/criger/stochastic-master-equations/'
    save_dir = master_dir + 'python/saved_traj/2014_07_07/'
    open_dir = master_dir + 'python_test/saved_pulses/'
    input_dict = prep_dict_indep()
    
    #5 arguments, start, stop, point, n_points, and y_offset
    #start, stop, point, n_points, y_offset = map(float, argv[1:])
    
    #4 Arguments for testing genuine optimal control pulses
    start, stop, point, n_points = map(float, argv[1:])

    #y_offset /= 10. #Nasty hack to allow integer loop in batch_script
    total_height = start + (stop - start) * point / n_points
    
    var_names = ['total_height']

    gammas = [0., 0.] #Dephasing and relaxation rates
    gamma_names = ['gamma_d', 'gamma_r']

    #pulse_f = lambda total_height: pulse_from_file(total_height, open_dir + 'pulse_10_2000_.pkl')
    pulse_f = lambda tvec, total_height: pulse_from_file_jr(tvec, total_height, open_dir + 'pulse_10_2000_.pkl')
    #pulse_f = lambda tvec, amplitude: sine_pulse(tvec, amplitude, 10.)
    
    #pulse_f = lambda tvec, amplitude: raised_sine_pulse(tvec, amplitude, y_offset, 10.)
    
    result = fid_append((total_height,), var_names, gammas, gamma_names, pulse_f, input_dict)
    #result['y_offset']=y_offset

    #phil = open(save_dir + "result_dict_" + str(total_height)+'_'+str(y_offset),'w')
    phil = open(save_dir + "result_dict_oct_" + str(total_height),'w')
    try:
        pickle.dump(result, phil)
    finally:
        phil.close()





















